// Copyright 2021-2025 Buf Technologies, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

import {
  type AnyDesc,
  type DescEnum,
  type DescExtension,
  type DescFile,
  type DescMessage,
  type DescService,
  protoInt64,
  ScalarType,
} from "@bufbuild/protobuf";
import type { ImportSymbol } from "./import-symbol.js";
import { createImportSymbol } from "./import-symbol.js";
import { makeImportPathRelative } from "./import-path.js";
import { createJsDocTextFromDesc, formatJsDocBlock } from "./jsdoc.js";
import type { ExportStatement, Printable } from "./printable.js";
import type { RuntimeImports } from "./runtime-imports.js";

/**
 * FileInfo represents an intermediate type using for transpiling TypeScript internally.
 */
export interface FileInfo {
  name: string;
  content: string;
  preamble?: string | undefined;
}

/**
 * Represents a JavaScript, TypeScript, or TypeScript declaration file.
 */
export interface GeneratedFile {
  /**
   * Create a standard preamble the includes comments at the top of the
   * protobuf source file (like a license header), as well as information
   * about the code generator and its version.
   *
   * The preamble is always placed at the very top of the generated file,
   * above import statements.
   *
   * A file with a preamble but no other content is still considered empty,
   * and will not be generated unless the plugin option keep_empty_files=true
   * is set.
   */
  preamble(file: DescFile): void;

  /**
   * Add a line of code to the file.
   *
   * - string: Prints the string verbatim.
   * - number or boolean: Prints a literal.
   * - bigint: Prints an expression using protoInt64.parse().
   * - Uint8Array: Prints an expression that re-created the array.
   * - ImportSymbol: Adds an import statement and prints the name of the symbol.
   * - DescMessage or DescEnum: Imports the type if necessary, and prints the name.
   */
  print(...printables: Printable[]): void;

  /**
   * Add a line of code to the file with tagged template literal and
   * an optional array of Printables.
   * See print(Printable[]) for behavior when printing Printable items.
   */
  print(fragments: TemplateStringsArray, ...printables: Printable[]): void;

  /**
   * Create a string literal.
   */
  string(string: string): Printable;

  /**
   * Create an array literal.
   */
  array(elements: Printable[]): Printable;

  /**
   * Create a JSDoc comment block with the given text. Line breaks and white-space
   * stay intact.
   */
  jsDoc(text: string, indentation?: string): Printable;

  /**
   * Create a JSDoc comment block for the given message, enumeration, or other
   * descriptor. The comment block will contain the original comments from the
   * protobuf source, and annotations such as `@generated from message MyMessage`.
   */
  jsDoc(schema: Exclude<AnyDesc, DescFile>, indentation?: string): Printable;

  /**
   * Create a printable export statement. For example:
   *
   * ```ts
   * f.print(f.export("abstract class", "MyClass"), " {}")
   * ```
   *
   * Will generate as:
   * ```ts
   * export abstract class MyClass {}
   * ```
   *
   * Using this method is preferred over a calling print() with a literal export
   * statement. If the plugin option `js_import_style=legacy_commonjs` is set,
   * exports will automatically be generated for CommonJS.
   */
  export(declaration: string, name: string): Printable;

  /**
   * Import a message or enumeration generated by protoc-gen-es.
   */
  importShape(schema: DescMessage | DescEnum): ImportSymbol;

  /**
   * Import a message or enumeration JSON type generated by protoc-gen-es.
   *
   * Note that protoc-gen-es only generates JSON types with the plugin option
   * `json_types=true`.
   */
  importJson(desc: DescMessage | DescEnum): ImportSymbol;

  /**
   * Import a message Valid type generated by protoc-gen-es.
   *
   * Note that protoc-gen-es only generates valid types with the plugin option
   * `valid_types`.
   */
  importValid(desc: DescMessage): ImportSymbol;

  /**
   * Import a descriptor generated by protoc-gen-es.
   */
  importSchema(
    schema: DescMessage | DescEnum | DescExtension | DescService | DescFile,
    typeOnly?: boolean,
  ): ImportSymbol;

  /**
   * Import any symbol from a file or package.
   *
   * The import path can point to a package, for example `@foo/bar/baz.js`, or
   * to a file, for example `./bar/baz.js`.
   *
   * Note that while paths to a file begin with a `./`, they must be
   * relative to the project root. The import path is automatically made
   * relative to the current file.
   */
  import(name: string, from: string, typeOnly?: boolean): ImportSymbol;

  /**
   * In case you need full control over exports and imports, use print() and
   * formulate your own imports and exports based on this property.
   *
   * With the plugin option `js_import_style=legacy_commonjs`, this property
   * reports "legacy_commonjs", but only if the current target is "js".
   * This matches the behavior of import(), which also only generates CommonJS
   * under this condition.
   */
  readonly jsImportStyle: "module" | "legacy_commonjs";

  /**
   * Provides some symbols from the runtime library @bufbuild/protobuf.
   */
  readonly runtime: RuntimeImports;
}

export interface GeneratedFileController extends GeneratedFile {
  getFileInfo(): FileInfo;
}

export type ResolveDescImportFn = (
  desc: DescMessage | DescEnum | DescExtension | DescService | DescFile,
  typeOnly?: boolean,
) => ImportSymbol;
export type ResolveShapeImportFn = (
  desc: DescMessage | DescEnum,
) => ImportSymbol;
export type ResolveJsonImportFn = (
  desc: DescMessage | DescEnum,
) => ImportSymbol;
export type ResolveValidImportFn = (desc: DescMessage) => ImportSymbol;
export type RewriteImportFn = (path: string) => string;
export type CreatePreambleFn = (descFile: DescFile) => string;

export function createGeneratedFile(
  name: string,
  importPath: string,
  jsImportStyle: "module" | "legacy_commonjs",
  rewriteImport: RewriteImportFn,
  resolveDescImport: ResolveDescImportFn,
  resolveShapeImport: ResolveShapeImportFn,
  resolveJsonImport: ResolveJsonImportFn,
  resolveValidImport: ResolveValidImportFn,
  createPreamble: CreatePreambleFn,
  runtime: RuntimeImports,
): GeneratedFileController {
  let preamble: string | undefined;
  const el: El[] = [];
  return {
    preamble(file) {
      preamble = createPreamble(file);
    },
    print(
      printableOrFragments?: Printable | TemplateStringsArray,
      ...rest: Printable[]
    ) {
      let printables: Printable[];
      if (
        printableOrFragments != null &&
        Object.prototype.hasOwnProperty.call(printableOrFragments, "raw")
      ) {
        // If called with a tagged template literal
        printables = buildPrintablesFromFragments(
          printableOrFragments as TemplateStringsArray,
          rest,
        );
      } else {
        // If called with just an array of Printables
        printables =
          printableOrFragments != null
            ? [printableOrFragments as Printable, ...rest]
            : rest;
      }
      printableToEl(
        {
          el,
          runtime,
          resolveDescImport,
          resolveShapeImport,
          resolveJsonImport,
          resolveValidImport,
        },
        printables,
      );
      el.push("\n");
    },
    export(declaration, name) {
      return {
        kind: "es_export_stmt",
        name,
        declaration,
      } satisfies ExportStatement;
    },
    string(string) {
      return {
        kind: "es_string",
        value: string,
      };
    },
    array(elements) {
      const p: Printable = [];
      p.push("[");
      for (const [index, element] of elements.entries()) {
        p.push(element);
        if (index < elements.length - 1) {
          p.push(", ");
        }
      }
      p.push("]");
      return p;
    },
    jsDoc(textOrSchema, indentation) {
      return {
        kind: "es_jsdoc",
        text:
          typeof textOrSchema == "string"
            ? textOrSchema
            : createJsDocTextFromDesc(textOrSchema),
        indentation,
      };
    },
    importSchema(schema, typeOnly = false) {
      return resolveDescImport(schema, typeOnly);
    },
    importShape(schema) {
      return resolveShapeImport(schema);
    },
    importJson(desc) {
      return resolveJsonImport(desc);
    },
    importValid(desc) {
      return resolveValidImport(desc);
    },
    import(name, from, typeOnly = false) {
      return createImportSymbol(name, from, typeOnly);
    },
    jsImportStyle,
    runtime,
    getFileInfo() {
      return {
        name,
        content: elToContent(
          el,
          importPath,
          rewriteImport,
          jsImportStyle == "legacy_commonjs",
        ),
        preamble,
      };
    },
  };
}

type El = string | ImportSymbol | ExportStatement;

function elToContent(
  el: El[],
  importerPath: string,
  rewriteImportPath: RewriteImportFn,
  legacyCommonJs: boolean,
): string {
  if (el.length == 0) {
    return "";
  }
  const c: string[] = [];
  if (legacyCommonJs) {
    c.push(`"use strict";\n`);
    c.push(`Object.defineProperty(exports, "__esModule", { value: true });\n`);
    c.push("\n");
  }
  const symbolToIdentifier = processImports(
    el,
    importerPath,
    rewriteImportPath,
    (typeOnly, from, names) => {
      if (legacyCommonJs) {
        const p = names.map(({ name, alias }) =>
          alias == undefined ? name : `${name}: ${alias}`,
        );
        const what = `{ ${p.join(", ")} }`;
        c.push(`const ${what} = require(${escapeString(from)});\n`);
      } else {
        const p = names.map(({ name, alias }) =>
          alias == undefined ? name : `${name} as ${alias}`,
        );
        const what = `{ ${p.join(", ")} }`;
        if (typeOnly) {
          c.push(`import type ${what} from ${escapeString(from)};\n`);
        } else {
          c.push(`import ${what} from ${escapeString(from)};\n`);
        }
      }
    },
  );
  if (c.length > 0) {
    c.push("\n");
  }
  const legacyCommonJsExportNames: string[] = [];
  for (const e of el) {
    if (typeof e == "string") {
      c.push(e);
    } else {
      switch (e.kind) {
        case "es_symbol": {
          const ident = symbolToIdentifier.get(e.id);
          if (ident != undefined) {
            c.push(ident);
          }
          break;
        }
        case "es_export_stmt":
          if (legacyCommonJs) {
            legacyCommonJsExportNames.push(e.name);
          } else {
            c.push("export ");
          }
          if (e.declaration !== undefined && e.declaration.length > 0) {
            c.push(e.declaration, " ");
          }
          c.push(e.name);
          break;
      }
    }
  }
  if (legacyCommonJs) {
    if (legacyCommonJsExportNames.length > 0) {
      c.push("\n");
    }
    for (const name of legacyCommonJsExportNames) {
      c.push("exports.", name, " = ", name, ";\n");
    }
  }
  return c.join("");
}

interface PrintableToElOpt {
  el: El[];
  runtime: RuntimeImports;
  resolveDescImport: ResolveDescImportFn;
  resolveShapeImport: ResolveShapeImportFn;
  resolveJsonImport: ResolveShapeImportFn;
  resolveValidImport: ResolveValidImportFn;
}
function printableToEl(opt: PrintableToElOpt, printables: Printable[]): void {
  const { el } = opt;
  for (const p of printables) {
    if (Array.isArray(p)) {
      printableToEl(opt, p);
    } else {
      switch (typeof p) {
        case "string":
          el.push(p);
          break;
        case "number":
          if (Number.isNaN(p)) {
            el.push("globalThis.NaN");
          } else if (p === Number.POSITIVE_INFINITY) {
            el.push("globalThis.Infinity");
          } else if (p === Number.NEGATIVE_INFINITY) {
            el.push("-globalThis.Infinity");
          } else {
            el.push(p.toString(10));
          }
          break;
        case "boolean":
          el.push(p.toString());
          break;
        case "bigint":
          if (p == protoInt64.zero) {
            // Loose comparison will match between 0n and 0.
            el.push(opt.runtime.protoInt64, ".zero");
          } else {
            el.push(
              opt.runtime.protoInt64,
              p > 0 ? ".uParse(" : ".parse(",
              escapeString(p.toString()),
              ")",
            );
          }
          break;
        case "object":
          if (p instanceof Uint8Array) {
            if (p.length === 0) {
              el.push("new Uint8Array(0)");
            } else {
              el.push("new Uint8Array([");
              const strings: string[] = [];
              for (const n of p) {
                strings.push(
                  "0x" + n.toString(16).toUpperCase().padStart(2, "0"),
                );
              }
              el.push(strings.join(", "));
              el.push("])");
            }
            break;
          }
          switch (p.kind) {
            case "es_symbol":
            case "es_export_stmt":
              el.push(p);
              break;
            case "es_desc_ref":
              el.push(opt.resolveDescImport(p.desc, p.typeOnly));
              break;
            case "es_shape_ref":
              el.push(opt.resolveShapeImport(p.desc));
              break;
            case "es_json_type_ref":
              el.push(opt.resolveJsonImport(p.desc));
              break;
            case "es_valid_type_ref":
              el.push(opt.resolveValidImport(p.desc));
              break;
            case "es_jsdoc":
              el.push(formatJsDocBlock(p.text, p.indentation));
              break;
            case "es_string":
              el.push(escapeString(p.value));
              break;
            case "es_proto_int64":
              if (p.longAsString) {
                el.push(escapeString(p.value.toString()));
              } else {
                if (p.value == protoInt64.zero) {
                  // Loose comparison will match between 0n and 0.
                  el.push(opt.runtime.protoInt64, ".zero");
                } else {
                  switch (p.type) {
                    case ScalarType.UINT64:
                    case ScalarType.FIXED64:
                      el.push(
                        opt.runtime.protoInt64,
                        ".uParse(",
                        escapeString(p.value.toString()),
                        ")",
                      );
                      break;
                    default:
                      el.push(
                        opt.runtime.protoInt64,
                        ".parse(",
                        escapeString(p.value.toString()),
                        ")",
                      );
                      break;
                  }
                }
              }
              break;
            default:
              throw `cannot print ${typeof p}`;
          }
          break;
        default:
          throw `cannot print ${typeof p}`;
      }
    }
  }
}

function buildPrintablesFromFragments(
  fragments: TemplateStringsArray,
  values: Printable[],
): Printable[] {
  const printables: Printable[] = [];
  fragments.forEach((fragment, i) => {
    printables.push(fragment);
    if (fragments.length - 1 !== i) {
      printables.push(values[i]);
    }
  });

  return printables;
}

type MakeImportStatementFn = (
  typeOnly: boolean,
  from: string,
  names: { name: string; alias?: string }[],
) => void;

function processImports(
  el: El[],
  importerPath: string,
  rewriteImportPath: RewriteImportFn,
  makeImportStatement: MakeImportStatementFn,
) {
  // identifiers to use in the output
  const symbolToIdentifier = new Map<string, string>();
  // symbols that need a value import (as opposed to a type-only import)
  const symbolToIsValue = new Map<string, true>();
  // taken in this file
  const identifiersTaken = new Set<string>();
  // foreign symbols need an import
  const foreignSymbols: ImportSymbol[] = [];

  // Walk through all symbols used and populate the collections above.
  for (const s of el) {
    if (typeof s != "object") {
      continue;
    }
    switch (s.kind) {
      case "es_symbol":
        symbolToIdentifier.set(s.id, s.name);
        if (!s.typeOnly) {
          // a symbol is only type-imported as long as all uses are type-only
          symbolToIsValue.set(s.id, true);
        }
        if (s.from === importerPath) {
          identifiersTaken.add(s.name);
        } else {
          foreignSymbols.push(s);
        }
        break;
      case "es_export_stmt":
        identifiersTaken.add(s.name);
        break;
    }
  }

  // Walk through all foreign symbols and make their identifiers unique.
  const handledSymbols = new Set<string>();
  for (const s of foreignSymbols) {
    if (handledSymbols.has(s.id)) {
      continue;
    }
    handledSymbols.add(s.id);
    if (!identifiersTaken.has(s.name)) {
      identifiersTaken.add(s.name);
      continue;
    }
    let i = 1;
    let alias: string;
    for (;;) {
      // We choose '$' because it is invalid in proto identifiers.
      alias = `${s.name}$${i}`;
      if (!identifiersTaken.has(alias)) {
        break;
      }
      i++;
    }
    identifiersTaken.add(alias);
    symbolToIdentifier.set(s.id, alias);
  }

  // Group foreign symbols (imports) by their source (from).
  type Imp = {
    types: Map<string, string | undefined>; // type-only imports, name to (optional) alias
    values: Map<string, string | undefined>; // value imports, name to (optional) alias
  };
  const sourceToImport = new Map<string, Imp>();
  for (const s of foreignSymbols) {
    let i = sourceToImport.get(s.from);
    if (i == undefined) {
      i = {
        types: new Map<string, string | undefined>(),
        values: new Map<string, string | undefined>(),
      };
      sourceToImport.set(s.from, i);
    }
    let alias = symbolToIdentifier.get(s.id);
    if (alias == s.name) {
      alias = undefined;
    }
    if (symbolToIsValue.get(s.id)) {
      i.values.set(s.name, alias);
    } else {
      i.types.set(s.name, alias);
    }
  }

  // Make import statements.
  const handledSource = new Set<string>();
  const buildNames = (map: Map<string, string | undefined>) => {
    const names: { name: string; alias?: string }[] = [];
    map.forEach((value, key) => names.push({ name: key, alias: value }));
    names.sort((a, b) => a.name.localeCompare(b.name));
    return names;
  };
  for (const s of foreignSymbols) {
    if (handledSource.has(s.from)) {
      continue;
    }
    handledSource.add(s.from);
    const i = sourceToImport.get(s.from);
    if (i == undefined) {
      // should never happen
      continue;
    }
    const from = makeImportPathRelative(
      importerPath,
      rewriteImportPath(s.from),
    );
    if (i.types.size > 0) {
      makeImportStatement(true, from, buildNames(i.types));
    }
    if (i.values.size > 0) {
      makeImportStatement(false, from, buildNames(i.values));
    }
  }

  return symbolToIdentifier;
}

function escapeString(value: string): string {
  return (
    '"' +
    value
      .split("\\")
      .join("\\\\")
      .split('"')
      .join('\\"')
      .split("\r")
      .join("\\r")
      .split("\n")
      .join("\\n") +
    '"'
  );
}
